
# Function to do MCMC sampling
MyCAR <- function(Y,adjs,X,N,firstday=1,nugget=T,spatial=T,temporal=T,
                  iters=200,burn=100,thin=1,update=1,initVals=NULL,
                  showProgress=0, dumpProgress=0) {

  library(spam)
  ptm <- proc.time()
  
  # Set up
  ns         <- nrow(Y) 
  ms         <- unlist(lapply(adjs,length))
  adj1       <- rep(1:ns,times=ms)
  adj2       <- unlist(adjs)
  
  # Set up time adjacency
  nt         <- ncol(Y)
  mt         <- c(1,rep(2,nt-2),1)
  adjt       <- list()
  adjt[[1]]  <- 2
  for(t in 2:(nt-1)){adjt[[t]]<-c(t-1,t+1)}
  adjt[[nt]] <- nt-1
  if(!is.matrix(N)){N <- matrix(N,ns,nt)}
  
  # Computation for the determinant
  ADJt  <- as.spam(as.matrix(dist(1:nt))==1)
  ADJs <- as.spam(matrix(0,ns,ns))
  for(s in 1:ns){ ADJs[s,adjs[[s]]] <-1 }
  Cs <- sweep(ADJs,1,sqrt(ms),"/")
  Cs <- sweep(Cs,2,sqrt(ms),"/")
  Ct <- sweep(ADJt,1,sqrt(mt),"/")
  Ct <- sweep(Ct,2,sqrt(mt),"/")
  ds <- eigen(Cs)$val
  dt <- eigen(Ct)$val
  rm(Cs,Ct)
  
  # Initial values
  # lambda(t) = N*[exp(e)+exp(v)]
  # v(t) ~iid N(muv, tauv)
  # e(t) ~indep N(theta(t), taue)
  # theta = mn_theta(alpha) + STCAR
  p     <- length(X)
  beta  <- rep(0,p)
  taus  <- 1
  taue  <- 10
  tauv  <- 10
  muv   <- -2
  rhos  <- ifelse(spatial,0.95,0)
  rhot  <- ifelse(temporal,0.50,0)
  e     <- log((Y+10)/(N+10))
  theta <- 0.95*e + 0.05*mean(e)
  v     <- ifelse(nugget,-2,-Inf)*matrix(1,ns,nt)
  if (!is.null(initVals)) {
    beta <- initVals[1:p]
    taue <- initVals[c("taue")]
    taus <- initVals[c("taus")]
    rhos <- initVals[c("rhos")]
    rhot <- initVals[c("rhot")]
    muv  <- initVals[c("muv")]
    tauv <- initVals[c("tauv")]
  }
  
  
  # Indictors of which observations should be used
  IN <- matrix(0,ns,nt)
  IN[,firstday:nt]<-1
  Y  <- IN*Y
  N  <- IN*N
  mu <- 0
  for (j in 1:p) { mu <- mu + X[[j]]*beta[j] }
  MH <- att <- acc <- c(2,2,0.05,0.05,0.05,0.05)
  M1 <- thetahat1 <- thetahat2 <- 0
  keepers <- matrix(NA,iters,p+6)
  dev     <- rep(NA,iters)
  colnames(keepers) <- c(names(X),"taue","taus","rhos","rhot","muv","tauv")
  
  tick <- proc.time()[3]
  for(iter in 1:iters){ if (iter%%100==0) { cat("\n",iter," ") } else { cat(".") }
    for(ttt in 1:thin){
      
      # Update nugget effect, e
      cane <- e+MH[1]*rnorm(ns*nt)/sqrt(taue)
      R    <- dpois(Y,N*(exp(cane)+exp(v)),log=TRUE) -  dpois(Y,N*(exp(e)+exp(v)),log=TRUE) +
              dnorm(cane,theta,1/sqrt(taue),log=TRUE) - dnorm(e,theta,1/sqrt(taue),log=TRUE)
      R[is.na(R)] <- -Inf
      accept      <- log(runif(ns*nt))<R
      e[accept]   <- cane[accept] 
      acc[1]      <- acc[1] + sum(accept)
      att[1]      <- att[1] + ns*nt
      
      if(nugget){
        canv <- v+MH[2]*rnorm(ns*nt)/sqrt(tauv)
        R    <- dpois(Y,N*(exp(e)+exp(canv)),log=TRUE) - dpois(Y,N*(exp(e)+exp(v)),log=TRUE) +
                dnorm(canv,muv,1/sqrt(tauv),log=TRUE)  - dnorm(v,muv,1/sqrt(tauv),log=TRUE)
        R[is.na(R)] <- -Inf
        accept      <- log(runif(ns*nt))<R
        v[accept]   <- canv[accept] 
        acc[2]      <- acc[2] + sum(accept)
        att[2]      <- att[2] + ns*nt
      } 
      
      taue <- rgamma(1,ns*nt/2+0.1,sum((e-theta)^2)/2+0.1)
      if(nugget){
        tauv <- rgamma(1,ns*nt/2+0.1,sum((v-muv)^2)/2+0.1)
        VVV  <- tauv*ns*nt+0.01
        MMM  <- tauv*sum(v)
        muv  <- rnorm(1,MMM/VVV,1/sqrt(VVV))
      }  
      
      # Update spatiotemporal process theta
      R   <- theta-mu
      Qt  <- taus*(diag(mt)-rhot*ADJt)
      for(j in unique(ms)){
        VVV <- solve(j*Qt + taue*diag(nt))
        PPP <- t(chol(VVV))
        for(k in which(ms==j)){
          MMM <- R[adjs[[k]],]
          if(j>1){MMM <- colSums(MMM)}
          MMM       <- Qt%*%(j*mu[k,]+rhos*MMM) + taue*e[k,]
          theta[k,] <- VVV%*%MMM + PPP%*%rnorm(nt)
          R[k,]     <- theta[k,]-mu[k,]
        }
      }
      SS   <- qf(R,rhos,rhot,ns,nt,ms,mt,adj1,adj2)
      taus <- rgamma(1,ns*nt/2+0.1,SS/2+0.1)
      
      # Do updates on the z scale
      if(spatial){
        att[5]  <- att[5] + 1
        z       <- qnorm(pbeta(rhos,1,1))
        canz    <- rnorm(1,z,MH[5])
        canrhos <- qbeta(pnorm(canz),1,1)
        canSS   <- qf(R,canrhos,rhot,ns,nt,ms,mt,adj1,adj2)
        MHR     <- 0.5*logdet(canrhos,rhot,ds,dt) - 0.5*taus*canSS -
                   0.5*logdet(rhos,rhot,ds,dt) + 0.5*taus*SS +
                   dnorm(canz,log=TRUE)-dnorm(z,log=TRUE)
        if(canrhos<0.995){if(log(runif(1))<MHR){
          rhos<-canrhos;SS<-canSS;acc[5] <- acc[5] + 1
        }}   
      }
      if(temporal){ 
        att[6]  <- att[6] + 1
        z       <- qnorm(pbeta(rhot,1,1))
        canz    <- rnorm(1,z,MH[6])
        canrhot <- qbeta(pnorm(canz),1,1)
        canSS   <- qf(R,rhos,canrhot,ns,nt,ms,mt,adj1,adj2)
        MHR     <- 0.5*logdet(rhos,canrhot,ds,dt) - 0.5*taus*canSS -
          0.5*logdet(rhos,rhot,ds,dt) + 0.5*taus*SS +
          dnorm(canz,log=TRUE)-dnorm(z,log=TRUE)
        if(log(runif(1))<MHR){
          rhot<-canrhot;SS<-canSS;acc[6] <- acc[6] + 1
        }   
      }
      
      # Update beta
      Qt   <- taus*(diag(mt)-rhot*ADJt)
      L    <- chol(Qt)
      for(j in 1:p){
        A       <- X[[j]]
        mu      <- mu - A*beta[j]
        La      <- L%*%t(A)
        Lt      <- L%*%t(theta-mu)
        VVV     <- 0.01 + sum(ms*colSums(La*La))-rhos*sum(La[,adj1]*La[,adj2])
        MMM     <- 0    + sum(ms*colSums(La*Lt))-rhos*sum(La[,adj1]*Lt[,adj2])
        beta[j] <- rnorm(1,MMM/VVV,1/sqrt(VVV))
        mu      <- mu + A*beta[j]
      }
    } # end thin
    
    # Tuning
    if (iter<burn){ for(j in 1:length(MH)){if(att[j]>25){
      if(acc[j]/att[j]<0.3){MH[j]<-MH[j]*0.8}
      if(acc[j]/att[j]>0.5){MH[j]<-MH[j]*1.2}
      acc[j]<-att[j]<-0
    }}}
    
    # Keep track of output
    keepers[iter,] <- c(beta,taue,taus,rhos,rhot,muv,tauv)
    MM             <- exp(e)+exp(v)
    dev[iter]      <- -2*sum(dpois(Y,N*MM,log=TRUE))
    if (iter>burn) {
      thetahat1 <- thetahat1+theta/(iters-burn)
      thetahat2 <- thetahat2+theta*theta/(iters-burn)
      M1        <- M1 + MM/(iters-burn)
    }
    
    # Evolving trace plots
    varsToPlot <- which(names(X) %in% c("A","Abar","Ahat","Ahatbar","Ahat2","Ahatbar2","Atime",
                    "popdens","Apopdens","AhatAhatbar","dayssincefirstcase","pm25",
                    "dayssincefirstcase2","time","time2",
                    "time1","time2","time3","baselineA","time1A","time2A","time3A"))
    if (length(varsToPlot)>9) {  
      Mfrow <- c(3,4)  
    } else if (length(varsToPlot)>6) {  
      Mfrow <- c(3,3)  
    } else {
      Mfrow <- c(2,3)  
    }
    if (showProgress & iter%%showProgress==0) {
      par(mfrow=Mfrow, oma=c(0,0,2,0), mar=2*c(1,1,1,1))
      for (j in varsToPlot) {
        plot(keepers[1:iter,j],main=colnames(keepers)[j],type="l")
        abline(v=burn); abline(h=0)
      }
    }

    
  } # end iters
  time <- (proc.time() - ptm)[3]/3600

  # Trace plots
  pdf(paste0("convergence.pdf"),width=8,height=8)
  par(mfrow=Mfrow, oma=c(0,0,2,0), mar=2*c(1,1,1,1))
  for (j in varsToPlot) {
    plot(keepers[1:iter,j],main=colnames(keepers)[j],type="l")
    abline(v=burn); abline(h=0)
  }
  dev.off()
  
  dbar <- mean(dev[burn:iters])
  dhat <- -2*sum(dpois(Y,N*M1,log=TRUE))
  pD   <- dbar-dhat
  DIC  <- dbar + pD
  tock <- proc.time()[3]
  out  <- list(keepers=keepers, theta_hat=thetahat1, theta_var=thetahat2-thetahat1^2,
               min=(tock-tick)/60, DIC=list(DIC=DIC,dbar=dbar,pD=pD), dev=dev)
  return(out)
}


# Det of the spatiotemporal CAR covariance
logdet <- function(rhos,rhot,ds,dt) {
  length(ds)*sum(log(1-rhot*dt)) + length(dt)*sum(log(1-rhos*ds))
}


# Compute CAR quadratic form
# t(R)%*%[SigmaS(rhos)\kroneckerSigmaT(rhot)]%*%R
qf <- function(R,rhos,rhot,ns,nt,ms,mt,adj1,adj2){
  SS <- ms%*%(R^2)%*%mt-rhos*sum((R[adj1,]*R[adj2,])%*%mt)
  SS <- SS-rhot*sum(ms%*%(R[,-1]*R[,-nt]))
  SS <- SS+rhot*rhos*sum(R[adj1,-1]*R[adj2,-nt])
  return(as.vector(SS))
}


# Compute the weighted function for the reporting lag
make.w <- function(nt,lag,bw){
  # w[j,t] = weight of lambda[j] on Z[t]
  w <- matrix(0,nt,nt)
  for(j in 1:nt){
    ww    <- (1:nt-j-lag)/bw
    w[j,] <- ifelse(1:nt<j,0,exp(-ww^2))
  } 
  return(w)
}


# lag stuff
lag_mat <- function(X,lag,nt){ 
  X[,(lag+1):nt] <- X[,(lag+1):nt-lag] 
  return(X)
}

